// hardware/rtl/matrix_mult.v
`timescale 1ns/1ps

module pe_cell #(
  parameter WIDTH = 8
)(
  input  wire             clk,
  input  wire             rst_n,
  input  wire [WIDTH-1:0] a_in,    // 横向流动数据
  input  wire [WIDTH-1:0] b_in,    // 纵向流动数据
  input  wire [15:0]      c_in,    // 累加输入
  output reg  [WIDTH-1:0] a_out,
  output reg  [WIDTH-1:0] b_out,
  output reg  [15:0]      c_out
);
  // 寄存器锁存输入数据
  always @(posedge clk or negedge rst_n) begin
    if (!rst_n) begin
      a_out <= 0;
      b_out <= 0;
    end else begin
      a_out <= a_in;
      b_out <= b_in;
    end
  end

  // 计算核心：c_out = a_in * b_in + c_in
  wire [15:0] mult_result;
  assign mult_result = $signed(a_in) * $signed(b_in);

  always @(posedge clk or negedge rst_n) begin
    if (!rst_n)
      c_out <= 0;
    else
      c_out <= mult_result + c_in;
  end
endmodule

module matrix_mult #(
  parameter SIZE  = 8,     // 阵列尺寸
  parameter WIDTH = 8      // 数据位宽
)(
  input  wire                     clk,
  input  wire                     rst_n,
  input  wire [WIDTH*SIZE-1:0]    a_matrix,  // 输入矩阵A (行优先)
  input  wire [WIDTH*SIZE-1:0]    b_matrix,  // 输入矩阵B (列优先)
  output wire [15:0]              result[SIZE][SIZE] // 输出矩阵
);

  // PE阵列间连接信号
  wire [WIDTH-1:0] a_bus [SIZE:0][SIZE:0];
  wire [WIDTH-1:0] b_bus [SIZE:0][SIZE:0];
  wire [15:0]      c_bus [SIZE:0][SIZE:0];

  // 初始化边界信号
  generate
    genvar i, j;
    for (i = 0; i <= SIZE; i = i + 1) begin
      assign a_bus[i][0] = 0;
      assign b_bus[0][i] = 0;
      assign c_bus[i][0] = 0;
      assign c_bus[0][i] = 0;
    end
  endgenerate

  // 输入数据加载
  generate
    for (i = 0; i < SIZE; i = i + 1) begin: load_a
      assign a_bus[i+1][0] = a_matrix[WIDTH*(SIZE-i)-1 -: WIDTH];
    end
    for (j = 0; j < SIZE; j = j + 1) begin: load_b
      assign b_bus[0][j+1] = b_matrix[WIDTH*(SIZE-j)-1 -: WIDTH];
    end
  endgenerate

  // PE阵列生成
  generate
    for (i = 0; i < SIZE; i = i + 1) begin: row
      for (j = 0; j < SIZE; j = j + 1) begin: col
        pe_cell #(
          .WIDTH(WIDTH)
        ) u_pe (
          .clk    (clk),
          .rst_n  (rst_n),
          .a_in   (a_bus[i][j]),
          .b_in   (b_bus[i][j]),
          .c_in   (c_bus[i][j]),
          .a_out  (a_bus[i][j+1]),
          .b_out  (b_bus[i+1][j]),
          .c_out  (c_bus[i+1][j+1])
        );

        // 结果捕获
        assign result[i][j] = c_bus[i+1][j+1];
      end
    end
  endgenerate

endmodule